import os
from os.path import join as osj
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

import faiss,pdb,torch
import gradio as gr
from transformers import AutoModel, AutoProcessor
from transformers.image_utils import load_image

SIGLIP_PATH = '/data/models/siglip2-so400m-patch16-512'

def get_siglip2():
    processor = AutoProcessor.from_pretrained(SIGLIP_PATH,)
    model = AutoModel.from_pretrained(SIGLIP_PATH, device_map="cuda").eval()
    return processor,model.cuda()
def extract_feature_siglip2( filepath='',file_pil=None , image_processor=None, model=None ):
    # prepare input image
    image = load_image(filepath) if file_pil is None else file_pil
    inputs = image_processor(images=[image], 
                             return_tensors="pt").to(model.device)

    with torch.no_grad():
        image_embeddings = model.get_image_features(**inputs)  # 1 1152
    
    return image_embeddings.cpu().numpy()

def search_by_faiss(query_emb, index, id_map , k=5):
    """ query_emb: 1x1154 tensor """
    D, I = index.search(query_emb, k)
    res_show = []
    for d,i in zip(D[0] , I[0]):
        # print('distance :' , d)
        # pdb.set_trace()
        imgpath = id_map[i]['name']
        box = id_map[i]['box']

        res_show.append( load_image(imgpath).crop(box) ) 
    result_str = "\n".join([
        f"结果 {idx+1}: 距离={d:.4f} | 图片={id_map[i]['name']} | 坐标={id_map[i]['box']}"
        for idx, (d, i) in enumerate(zip(D[0], I[0]))
    ])
    return res_show,[{"distance": d, **id_map[i]} for d,i in zip(D[0],I[0])],result_str

# 检索
def load_faiss_index(index_path: str, gpu_id: int = -1) -> faiss.Index:
    """
    加载FAISS索引文件，支持自动切换CPU/GPU
    
    参数:
        index_path: .faiss文件路径
        gpu_id: 使用的GPU ID (-1表示CPU)
    
    返回:
        faiss.Index对象
    
    异常:
        FileNotFoundError: 文件不存在时抛出
        RuntimeError: 索引加载失败时抛出
    """
    if not os.path.exists(index_path):
        raise FileNotFoundError(f"FAISS索引文件不存在: {index_path}")
    
    try:
        # 基础加载
        index = faiss.read_index(index_path)
        
        # GPU加速 (可选)
        if gpu_id >= 0:
            res = faiss.StandardGpuResources()
            index = faiss.index_cpu_to_gpu(res, gpu_id, index)
            print(f"已启用GPU加速 (Device {gpu_id})")
            
        return index
        
    except Exception as e:
        raise RuntimeError(f"加载FAISS索引失败: {str(e)}") from e
    
def load_idmap(idmap_path: str, verbose: bool = True) -> dict:
    """
    加载ID映射表
    
    参数:
        idmap_path: .pt文件路径
        verbose: 是否打印加载信息
    
    返回:
        {id: metadata} 格式的字典
    
    异常:
        FileNotFoundError: 文件不存在时抛出
        RuntimeError: 文件损坏时抛出
    """
    if not os.path.exists(idmap_path):
        raise FileNotFoundError(f"ID映射文件不存在: {idmap_path}")
    
    try:
        id_map = torch.load(idmap_path)
        if verbose:
            print(f"成功加载ID映射表，共 {len(id_map)} 条记录")
            
        # 兼容性检查
        if not isinstance(id_map, dict):
            raise ValueError("ID映射文件格式错误，应为字典类型")
            
        return id_map
        
    except Exception as e:
        raise RuntimeError(f"加载ID映射表失败: {str(e)}") from e
    

# 编写一个函数
def get_rag_by_faiss( imgpath , processor , model , cloth_type ,top_k=5 ):
    assert cloth_type in ['collar','sleeve','pockets'],"cloth type must in ['collar','sleeve','pockets']"
    t = cloth_type
    
    get_faiss_path = lambda t : f'/mnt/nas/shengjie/datasets/faiss_emb_{t}.faiss'
    get_idmap_path = lambda t : f'/mnt/nas/shengjie/datasets/faiss_emb_{t}_idmap.pt'
    # processor , model = get_siglip2()
    # 选择一个faiss和id map文件
    index = load_faiss_index( get_faiss_path(t) )
    id_map = load_idmap( get_idmap_path(t) )

    query_emb = extract_feature_siglip2( imgpath , 
                                        image_processor=processor,
                                        model=model )  # numpy
                

    # search
    # 确保输入的emb是numpy
    rag_images,rag_info,rag_info_str = search_by_faiss( query_emb , index , id_map , k=5 )
    return rag_images,rag_info_str



#################### Gradio界面 ######################
processor, model = None,None
get_img_dir = lambda t: f'/mnt/nas/shengjie/datasets/cloth_{t}_localimg/'

def gradio_search(input_img, cloth_type, top_k):
    global processor,model
    if processor is None or model is None:
        processor, model = get_siglip2()

    input_img_path = "./tmp.jpg"
    input_img.save(input_img_path)
    
    images, info_str = get_rag_by_faiss(input_img_path, processor, model, cloth_type, top_k)
    
    # 格式化输出
    gallery = images if images else [None] * top_k
    text_output = info_str if info_str else "无结果"
    print(text_output)
    return gallery, text_output

# 创建界面
def test_gradio():
    with gr.Blocks(title="服装部件检索系统") as demo:

        gr.Markdown("## 服装部件相似检索 (基于SigLIP+FAISS)")
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="上传查询图片")
                cloth_type = gr.Dropdown(
                    choices=['collar', 'sleeve', 'pockets'],
                    value='collar',
                    label="选择服装部件类型"
                )
                top_k = gr.Slider(1, 10, value=5, step=1, label="返回结果数量")
                submit_btn = gr.Button("检索", variant="primary")
            
            with gr.Column():
                gallery = gr.Gallery(
                    label="相似结果",
                    columns=5,
                    object_fit="contain",
                    height="auto"
                )
                similarity_text = gr.Textbox(label="相似度详情", interactive=False)
        
        submit_btn.click(
            fn=gradio_search,
            inputs=[input_image, cloth_type, top_k],
            outputs=[gallery, similarity_text]
        )
        

        # 示例图片
        gr.Examples(
            examples=[
                [os.path.join(get_img_dir('collar'), "collar_0000916.jpg"), "collar", 5],
                [os.path.join(get_img_dir('sleeve'), "sleeve_0000061.jpg"), "sleeve", 4],
                [os.path.join(get_img_dir('pockets'), "pockets_20250625_Cargo-Pocket_244.jpg"), "pockets", 6]
            ],
            inputs=[input_image, cloth_type, top_k],
            outputs=[gallery, similarity_text],
            fn=gradio_search,
            cache_examples=False
        )
    # 启动界面
    demo.launch(
        server_name="0.0.0.0",
        server_port=20018,
        share=False
    )



##################  测试 函数 #######################
def test():
    processor , model = get_siglip2()

    types = ['collar','sleeve','pockets']
    for t in types:
        img_dir = f'/mnt/nas/shengjie/datasets/cloth_{t}_localimg/'
        for entry in os.scandir( img_dir ):
            if entry.is_file() and not entry.name.endswith('.txt'):
                imgpath = osj( img_dir , entry.name )

                rag_images,rag_info = \
                        get_rag_by_faiss( imgpath,processor,model,
                                            t,top_k=5 )

            break

if __name__=='__main__':
    # test()

    test_gradio()
